from PIL import Image
from torchvision import transforms
from torch.utils.data import Dataset
import os


class DataSetRace(Dataset):
    def __init__(self,dir,size=256,default_label=None) -> None:
        super(DataSetRace,self).__init__()
        
        self.size = size 
        self.data = self.get_images_by_oswalk(dir)
        self.T = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))
        ])
        self.default_label = default_label
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, index):
        image_path = self.data[index]
        image = Image.open(image_path).convert('RGB')
        image = image.resize((self.size,self.size))
        
        file_name= image_path.split('\\')[-1]
        if self.default_label is None:
            race = int(file_name.split('_')[2])
            if race == 1:
                race = 1
            else:
                race = 0
        else:
            race = self.default_label
        
        return self.T(image),race
    @staticmethod
    def get_images_by_oswalk(dir):
        data = []
        for root,dir,files in os.walk(dir):
            for file in files:
                if file.endswith(('jpg','png','jpeg','bmp')):
                    data.append(os.path.join(root,file))
        return data